import os
import warnings
warnings.filterwarnings('ignore')


os.environ["OPENAI_API_KEY"] = 'sk-f2M0jWxDEb7IJKiuCnMsvKE7TLSYCA36pyl8gWA1XLaHnD8l'
os.environ["OPENAI_BASE_URL"] = 'https://api.f2gpt.com/v1'
from dotenv import load_dotenv  # 用于加载环境变量

load_dotenv()  # 加载 .env 文件中的环境变量

# 数据的准备和载入

# 1.Load 导入Document Loaders
from langchain.document_loaders import PyPDFLoader
from langchain.document_loaders import Docx2txtLoader
from langchain.document_loaders import TextLoader

# 加载Documents
# 文档的存放目录
base_dir = "./OneFlower"
documents = []
for file in os.listdir(base_dir):
    # 构建完整的文件路径
    file_path = os.path.join(base_dir, file)
    if file.endswith(".pdf"):
        loader = PyPDFLoader(file_path)
        documents.extend(loader.load())
    elif file.endswith(".docx"):
        loader = Docx2txtLoader(file_path)
        documents.extend(loader.load())
    elif file.endswith(".txt"):
        loader = TextLoader(file_path)
        documents.extend(loader.load())

# 文本的分割
# 2.Split 将Documents切分成块以便后续进行嵌入和向量存储
from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=10)
chunked_documents = text_splitter.split_documents(documents)

# 向量数据库存储
from langchain.vectorstores import Qdrant
from langchain.embeddings import OpenAIEmbeddings

vectorstore = Qdrant.from_documents(
    documents=chunked_documents,  # 以分块的文档
    embedding=OpenAIEmbeddings(),  # 用OpenAI的Embedding Model做嵌入
    location=":memory:",  # in-memory 存储
    collection_name="my_documents",  # 指定collection_name
)

# 相关信息的获取
# 4. Retrieval 准备模型和Retrieval链
import logging  # 导入logging模块
from langchain.chat_models import ChatOpenAI  # 导入ChatOpenAI模型
from langchain.retrievers.multi_query import MultiQueryRetriever  # 导入MultiQueryRetriever
from langchain.chains import RetrievalQA  # 导入RetrievalQA链

# 设置logging
logging.basicConfig()
logging.getLogger("langchain.retrievers.multi_query").setLevel(logging.INFO)

# 实例化一个大模型工具 - OpenAI的GPT-3.5
llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)

# 实例化一个MultiQueryRetriever
retriever_from_llm = MultiQueryRetriever.from_llm(
    retriever=vectorstore.as_retriever(),  # 检索器
    llm=llm,  # 使用的语言模型
)

# 实例化一个RetrievalQA链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,  # 使用的语言模型
    retriever=retriever_from_llm,  # 检索器
)

# 生成回答并展示

# 5.Output 问答系统的UI实现
from flask import Flask, request, render_template

app = Flask(__name__)  # Flask APP


@app.route('/', methods=['GET', 'POST'])
def home():
    if request.method == 'POST':
        # 接收用户输入作为问题
        question = request.form.get('question')

        # RetrievalQA链 - 读入问题，生成答案
        result = qa_chain({"query": question})

        # 把大模型的回答结果返回网页进行渲染
        return render_template('index.html', result=result)

    return render_template('index.html')


if __name__ == "__main__":
    app.run(host='127.0.0.1', debug=True, port=5000)